[tests] Split Pipeline Group Offloading Block-Level and Leaf-Level Tests by dg845 · Pull Request #14635 · huggingface/diffusers

dg845 · 2026-08-28T06:10:58Z

What does this PR do?

This PR splits the GroupOffloadTesterMixin.test_group_offloading_inference test into separate block-level and leaf-level tests:

test_group_offloading_inference_block_level
test_group_offloading_inference_leaf_level

with the motivation that some components tend be block-level offloadable while not being leaf-level offloadable. To help support this, the PR makes the current hardcoded component list used by that test configurable via BasePipelineTesterConfig:

the PR default behavior is to test offloading on all components, except those which are excluded for that offloading level
Modules can be excluded in several ways:
- group_offloading_leaf_level_exclude_modules: excludes modules from leaf-level offloading tests, added previously in [tests] Add Config to Exclude Modules from Leaf-Level Group Offloading #14564
- group_offloading_exclude_modules: excludes modules from both leaf-level and block-level offloading tests
- group_offloading_onload_component_names: includes modules which should be kept on the accelerator rather than offloaded; the motivating case is vae components, which do not offload correctly if tiling is enabled

The new changes are then applied to the Ideogram 4 tests (the motivating example) and the LTX-2 tests (which overrides test_group_offloading_inference, so needs to be changed due to that test name no longer existing on the PR).

Fixes #14618

Before submitting

Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

@sayakpaul
@DN6

…ponents `GroupOffloadTesterMixin.test_group_offloading_inference` ran both offload levels in one test body, so a pipeline that failed at one level had to skip both. Split it into `test_group_offloading_inference_block_level` and `test_group_offloading_inference_leaf_level`, sharing the helpers the test body used to define inline, and compare against the class-scoped `base_pipe_output` rather than rebuilding a baseline per level. The set of components to offload was a hardcoded list of eight names, so a pipeline with a component under any other name had it silently left on CPU for the forward pass to trip over. Derive the set instead: every `nn.Module` component is offloaded unless the config lists it in `group_offloading_leaf_level_exclude_modules`, the new `group_offloading_exclude_modules`, or `group_offloading_onload_component_names`. A name in either exclusion list that matches no component on the pipeline fails as the typo it is. Ideogram4's `unconditional_transformer` was one of the silently dropped components, which is why its group offload test was skipped; it now needs no declaration at all, and `text_encoder` picks up block-level coverage it never had. LTX2's skip goes the same way, with `audio_vae` declared alongside the other VAEs the tests keep on the accelerator. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

sayakpaul · 2026-08-28T06:14:17Z

Do we want to hold off a bit until all the pipelines are migrated to use the new mixins? This way, it's easier to do this all at once?

dg845 · 2026-08-28T06:24:22Z

Self-Review Report

Self-review — tests/group-offloading-separate-levels (round 3)

git diff main...HEAD — one commit, c2466f611, 5 files, +118/−111. Rubric: .ai/references/review-rules.md + testing.md.

Blocking

The branch no longer merges cleanly into main. tests/pipelines/ideogram4/test_pipeline_ideogram4.py

A test-merge (git merge-tree $(git merge-base main <branch>) main <branch>) produces one conflict hunk. The branch deletes the group-offload skip at main:310-318; 741c7d0c6 added TestIdeogram4PipelineLoRA and TestIdeogram4PipelineLoRAMemory immediately after it, so git can't place the deletion:

   def test_sequential_offload_forward_pass_twice(self, expected_max_diff=2e-4):
       pass

+<<<<<<< .our
@pytest.mark.skip(reason=("Block-level group offloading cannot cover text_encoder: ...
class TestIdeogram4PipelineLoRAMemory(Ideogram4PipelineTesterConfig, LoraMemoryTesterMixin):
+=======
+>>>>>>> .their

The other four files merge clean. Resolution is mechanical — keep main's two new LoRA classes, drop the skip block — but the branch needs a rebase onto a8c33fe99 and a re-run afterwards, since every result in this report was measured against the old base. This is the one thing that must happen before the PR is openable.

Non-blocking

Two new mixin consumers on main, both verified harmless — but unverified by execution.

TestLumina2PipelineMemory composes MemoryTesterMixin and doesn't override the renamed test, so the split reaches it as two tests. Its get_dummy_components returns transformer, vae, scheduler, text_encoder, tokenizer — all inside the old hardcoded list ∪ onload list, so the derived-set change is a no-op for it.
TestIdeogram4PipelineLoRAMemory / TestLumina2PipelineLoRAMemory use LoraMemoryTesterMixin, which subclasses BaseLoraTesterMixin — not GroupOffloadTesterMixin — so the rename doesn't touch their parametrized ..._denoiser tests.

group_offloading_exclude_modules still has no consumer. common.py:82 — carried from rounds 1–2, verified functional at all three levels. Land now or with its first user; maintainer's call.

Summary

The change itself is in good shape — rounds 1 and 2 findings are all fixed and re-verified, the accounting closes with no unexplained test-state changes, and the nine sweep failures are confirmed pre-existing against a clean tree.

What changed since round 2 is entirely external: main moved, and it moved into the same file. The finding that matters is (1), and it's a rebase rather than a redesign — but every number in the table above predates it, so they're evidence about the old base, not a green light for the merged result.

Verdict: NEEDS CHANGES — solely on (1).

Fix before submitting — rebase onto a8c33fe99, resolve the ideogram4 hunk (keep main's LoRA classes, drop the skip), then re-run pytest tests/pipelines -k group_offloading plus the ideogram4 and lumina2 files to confirm (2).
Leave for the actual review — (3).

dg845 · 2026-08-28T06:31:52Z

Do we want to hold off a bit until all the pipelines are migrated to use the new mixins? This way, it's easier to do this all at once?

My current plan is to just have the core changes + Ideogram 4 changes (as proof of concept) + LTX-2 changes (since the CI will be red otherwise), and hold off on changes to other pipelines for now. Would that work?

sayakpaul · 2026-08-28T06:49:08Z

and hold off on changes to other pipelines for now. Would that work?

But this should be propagated to most pipelines, right?

dg845 · 2026-08-28T07:01:14Z

Right, I will follow up for the other pipelines either on this PR or a follow-up PR after the mixin migration is finished (or sufficiently stable).

sayakpaul

Thanks. The split is nice.

sayakpaul · 2026-08-28T11:27:44Z

  - In the first pass, just add tests related to `PipelineTesterMixin` and `MemoryTesterMixin`.
- **Declare a component that can't be offloaded — don't hand-write a skip.** Leaf-level offloading hooks only the supported leaf types (`nn.Linear`, `nn.Conv*`, `nn.Embedding` — see `_GO_LC_SUPPORTED_PYTORCH_LAYERS` in `src/diffusers/hooks/_common.py`) and onloads each on its own `forward`, so any code that reads a leaf's `.weight` instead of calling the leaf bypasses that leaf's hook and computes against offloaded weights. Which fix applies depends on who owns the component. For a diffusers model, set `_supports_group_offloading = False` on the `ModelMixin` subclass (as `HunyuanDiT2DModel` does) — both offload mixins honor the flag and skip themselves, so the gap is declared on the model instead of buried in a test file. For a third-party component you can't annotate, such as a `transformers` encoder, list it in `group_offloading_leaf_level_exclude_modules` on the config class; `enable_group_offload` keeps excluded components on the accelerator, so every other component stays covered — including the VAE, which the component-scoped `test_group_offloading_inference` leaves out. Block-level offloading is usually unaffected, hence the level in the name — a component that fails at both levels does need a skip.
+- **Declare a component that can't be offloaded — don't hand-write a skip.** Leaf-level offloading hooks only the supported leaf types (`nn.Linear`, `nn.Conv*`, `nn.Embedding` — see `_GO_LC_SUPPORTED_PYTORCH_LAYERS` in `src/diffusers/hooks/_common.py`) and onloads each on its own `forward`, so any code that reads a leaf's `.weight` instead of calling the leaf bypasses that leaf's hook and computes against offloaded weights. Which fix applies depends on who owns the component. For a diffusers model, set `_supports_group_offloading = False` on the `ModelMixin` subclass (as `HunyuanDiT2DModel` does) — both offload mixins honor the flag and skip themselves, so the gap is declared on the model instead of buried in a test file. For a third-party component you can't annotate, such as a `transformers` encoder, list it in `group_offloading_leaf_level_exclude_modules` on the config class; excluded components are kept on the accelerator, so every other component stays covered. Block-level offloading is usually unaffected, hence the level in the name, and `test_group_offloading_inference_block_level` still covers the component — a component that fails at both levels goes in `group_offloading_exclude_modules` instead, with a comment saying why.
+  - **Every `nn.Module` component is group offloaded unless a config list names it** — `group_offloading_leaf_level_exclude_modules`, `group_offloading_exclude_modules`, or `group_offloading_onload_component_names` (the VAE and friends, kept on the accelerator because tiling breaks stream tracing). A pipeline that adds a second denoiser or an extra encoder therefore gets it exercised without touching the shared mixin, and losing coverage takes naming the component. A name in an exclusion list that matches no component on the pipeline fails the test as a typo.


Let's provide an example here for easier navigation.

sayakpaul · 2026-08-28T11:28:23Z

    output_shape = (3, 16, 16)
-    # `encode_prompt` drives the Qwen3-VL decoder layers directly instead of calling `text_encoder.forward`, so the
-    # offloading hooks would leave its inputs on the offload device. Keep the text encoder out of group offloading.
+    # `encode_prompt` drives the Qwen3-VL decoder layers directly instead of calling `text_encoder.forward`, and


What is driving here? I think it's better to clarify that.

sayakpaul · 2026-08-28T11:28:50Z

+    group_offloading_onload_component_names = [
+        *BasePipelineTesterConfig.group_offloading_onload_component_names,
+        "audio_vae",
+    ]


I would prefer explicitly defining them.

sayakpaul · 2026-08-28T11:29:26Z

+    # exists: some tests enable tiling, and when accelerator streams are used the execution order of a tiled
+    # forward pass is not traced correctly, which errors out. Group offloading a VAE wants a warmup forward pass
+    # first (even on dummy inputs).
+    group_offloading_onload_component_names = ["vae", "vqvae", "image_encoder"]


We should keep them empty IMO to have users / agents explicitly set them.

sayakpaul · 2026-08-28T11:32:13Z

+        device mismatch. The onload names are not checked: they are a shared default covering several pipelines,
+        most of which have only some of them.
+        """
+        exclude = set(self.group_offloading_exclude_modules)


The scope of self.group_offloading_exclude_modules and this method are kind of intertwined. Can we split the method into two?

One that prepares exclude

One that takes exclude and pipe and raises if needed.

github-actions Bot added fixes-issue size/L PR with diff > 200 LOC tests and removed fixes-issue labels Aug 28, 2026

Merge branch 'main' into tests/group-offloading-separate-levels

b51d341

github-actions Bot added the fixes-issue label Aug 28, 2026

sayakpaul reviewed Aug 28, 2026

View reviewed changes

Navigation Menu

Uh oh!

[tests] Split Pipeline Group Offloading Block-Level and Leaf-Level Tests - #14635

[tests] Split Pipeline Group Offloading Block-Level and Leaf-Level Tests#14635
dg845 wants to merge 2 commits into
mainfrom
tests/group-offloading-separate-levels

dg845 commented Aug 28, 2026

Uh oh!

sayakpaul commented Aug 28, 2026

Uh oh!

dg845 commented Aug 28, 2026

Blocking

Non-blocking

Summary

Uh oh!

dg845 commented Aug 28, 2026

Uh oh!

sayakpaul commented Aug 28, 2026

Uh oh!

dg845 commented Aug 28, 2026

Uh oh!

sayakpaul left a comment

Uh oh!

sayakpaul Aug 28, 2026

Uh oh!

sayakpaul Aug 28, 2026

Uh oh!

sayakpaul Aug 28, 2026

Uh oh!

sayakpaul Aug 28, 2026

Uh oh!

sayakpaul Aug 28, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Uh oh!

Conversation

dg845 commented Aug 28, 2026

What does this PR do?

Before submitting

Who can review?

Uh oh!

sayakpaul commented Aug 28, 2026

Uh oh!

dg845 commented Aug 28, 2026

Blocking

Non-blocking

Summary

Uh oh!

dg845 commented Aug 28, 2026

Uh oh!

sayakpaul commented Aug 28, 2026

Uh oh!

dg845 commented Aug 28, 2026

Uh oh!

sayakpaul left a comment

Choose a reason for hiding this comment

Uh oh!

sayakpaul Aug 28, 2026

Choose a reason for hiding this comment

Uh oh!

sayakpaul Aug 28, 2026

Choose a reason for hiding this comment

Uh oh!

sayakpaul Aug 28, 2026

Choose a reason for hiding this comment

Uh oh!

sayakpaul Aug 28, 2026

Choose a reason for hiding this comment

Uh oh!

sayakpaul Aug 28, 2026

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants